{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1902a2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install langchain  langgraph  langchain-community"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c301f20",
   "metadata": {},
   "source": [
    "一. Build an agent with minimal code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "382c0c6c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\a\\.conda\\envs\\langchain_learn\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from langchain.chat_models import init_chat_model\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "\n",
    "llm = init_chat_model(\"deepseek:deepseek-chat\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fb209f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests, pathlib\n",
    "\n",
    "url = \"https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db\"\n",
    "local_path = pathlib.Path(\"Chinook.db\")\n",
    "\n",
    "if local_path.exists():\n",
    "    print(f\"{local_path} already exists, skipping download.\")\n",
    "else:\n",
    "    response = requests.get(url)\n",
    "    if response.status_code == 200:\n",
    "        local_path.write_bytes(response.content)\n",
    "        print(f\"File downloaded and saved as {local_path}\")\n",
    "    else:\n",
    "        print(f\"Failed to download the file. Status code: {response.status_code}\")\n",
    "\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "\n",
    "db = SQLDatabase.from_uri(\"sqlite:///Chinook.db\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b513feae",
   "metadata": {},
   "source": [
    "<!-- 建表语句 -->\n",
    "\n",
    "CREATE TABLE \"Album\" (\n",
    "\t\"AlbumId\" INTEGER NOT NULL,\n",
    "\t\"Title\" NVARCHAR(160) NOT NULL,\n",
    "\t\"ArtistId\" INTEGER NOT NULL,\n",
    "\tPRIMARY KEY (\"AlbumId\"),\n",
    "\tFOREIGN KEY(\"ArtistId\") REFERENCES \"Artist\" (\"ArtistId\")\n",
    ")\n",
    "\n",
    "/*\n",
    "3 rows from Album table:\n",
    "AlbumId\tTitle\tArtistId\n",
    "1\tFor Those About To Rock We Salute You\t1\n",
    "2\tBalls to the Wall\t2\n",
    "3\tRestless and Wild\t2\n",
    "*/\n",
    "\n",
    "\n",
    "CREATE TABLE \"Artist\" (\n",
    "\t\"ArtistId\" INTEGER NOT NULL,\n",
    "\t\"Name\" NVARCHAR(120),\n",
    "\tPRIMARY KEY (\"ArtistId\")\n",
    ")\n",
    "\n",
    "/*\n",
    "3 rows from Artist table:\n",
    "ArtistId\tName\n",
    "1\tAC/DC\n",
    "2\tAccept\n",
    "3\tAerosmith\n",
    "*/\n",
    "\n",
    "\n",
    "CREATE TABLE \"Customer\" (\n",
    "\t\"CustomerId\" INTEGER NOT NULL,\n",
    "\t\"FirstName\" NVARCHAR(40) NOT NULL,\n",
    "\t\"LastName\" NVARCHAR(20) NOT NULL,\n",
    "\t\"Company\" NVARCHAR(80),\n",
    "\t\"Address\" NVARCHAR(70),\n",
    "\t\"City\" NVARCHAR(40),\n",
    "\t\"State\" NVARCHAR(40),\n",
    "\t\"Country\" NVARCHAR(40),\n",
    "\t\"PostalCode\" NVARCHAR(10),\n",
    "\t\"Phone\" NVARCHAR(24),\n",
    "\t\"Fax\" NVARCHAR(24),\n",
    "\t\"Email\" NVARCHAR(60) NOT NULL,\n",
    "\t\"SupportRepId\" INTEGER,\n",
    "\tPRIMARY KEY (\"CustomerId\"),\n",
    "\tFOREIGN KEY(\"SupportRepId\") REFERENCES \"Employee\" (\"EmployeeId\")\n",
    ")\n",
    "\n",
    "/*\n",
    "3 rows from Customer table:\n",
    "CustomerId\tFirstName\tLastName\tCompany\tAddress\tCity\tState\tCountry\tPostalCode\tPhone\tFax\tEmail\tSupportRepId\n",
    "1\tLuís\tGonçalves\tEmbraer - Empresa Brasileira de Aeronáutica S.A.\tAv. Brigadeiro Faria Lima, 2170\tSão José dos Campos\tSP\tBrazil\t12227-000\t+55 (12) 3923-5555\t+55 (12) 3923-5566\tluisg@embraer.com.br\t3\n",
    "2\tLeonie\tKöhler\tNone\tTheodor-Heuss-Straße 34\tStuttgart\tNone\tGermany\t70174\t+49 0711 2842222\tNone\tleonekohler@surfeu.de\t5\n",
    "3\tFrançois\tTremblay\tNone\t1498 rue Bélanger\tMontréal\tQC\tCanada\tH2G 1A7\t+1 (514) 721-4711\tNone\tftremblay@gmail.com\t3\n",
    "*/\n",
    "\n",
    "\n",
    "CREATE TABLE \"Employee\" (\n",
    "\t\"EmployeeId\" INTEGER NOT NULL,\n",
    "\t\"LastName\" NVARCHAR(20) NOT NULL,\n",
    "\t\"FirstName\" NVARCHAR(20) NOT NULL,\n",
    "\t\"Title\" NVARCHAR(30),\n",
    "\t\"ReportsTo\" INTEGER,\n",
    "\t\"BirthDate\" DATETIME,\n",
    "\t\"HireDate\" DATETIME,\n",
    "\t\"Address\" NVARCHAR(70),\n",
    "\t\"City\" NVARCHAR(40),\n",
    "\t\"State\" NVARCHAR(40),\n",
    "\t\"Country\" NVARCHAR(40),\n",
    "\t\"PostalCode\" NVARCHAR(10),\n",
    "\t\"Phone\" NVARCHAR(24),\n",
    "\t\"Fax\" NVARCHAR(24),\n",
    "\t\"Email\" NVARCHAR(60),\n",
    "\tPRIMARY KEY (\"EmployeeId\"),\n",
    "\tFOREIGN KEY(\"ReportsTo\") REFERENCES \"Employee\" (\"EmployeeId\")\n",
    ")\n",
    "\n",
    "/*\n",
    "3 rows from Employee table:\n",
    "EmployeeId\tLastName\tFirstName\tTitle\tReportsTo\tBirthDate\tHireDate\tAddress\tCity\tState\tCountry\tPostalCode\tPhone\tFax\tEmail\n",
    "1\tAdams\tAndrew\tGeneral Manager\tNone\t1962-02-18 00:00:00\t2002-08-14 00:00:00\t11120 Jasper Ave NW\tEdmonton\tAB\tCanada\tT5K 2N1\t+1 (780) 428-9482\t+1 (780) 428-3457\tandrew@chinookcorp.com\n",
    "2\tEdwards\tNancy\tSales Manager\t1\t1958-12-08 00:00:00\t2002-05-01 00:00:00\t825 8 Ave SW\tCalgary\tAB\tCanada\tT2P 2T3\t+1 (403) 262-3443\t+1 (403) 262-3322\tnancy@chinookcorp.com\n",
    "3\tPeacock\tJane\tSales Support Agent\t2\t1973-08-29 00:00:00\t2002-04-01 00:00:00\t1111 6 Ave SW\tCalgary\tAB\tCanada\tT2P 5M5\t+1 (403) 262-3443\t+1 (403) 262-6712\tjane@chinookcorp.com\n",
    "*/\n",
    "\n",
    "\n",
    "CREATE TABLE \"Genre\" (\n",
    "\t\"GenreId\" INTEGER NOT NULL,\n",
    "\t\"Name\" NVARCHAR(120),\n",
    "\tPRIMARY KEY (\"GenreId\")\n",
    ")\n",
    "\n",
    "/*\n",
    "3 rows from Genre table:\n",
    "GenreId\tName\n",
    "1\tRock\n",
    "2\tJazz\n",
    "3\tMetal\n",
    "*/\n",
    "\n",
    "\n",
    "CREATE TABLE \"Invoice\" (\n",
    "\t\"InvoiceId\" INTEGER NOT NULL,\n",
    "\t\"CustomerId\" INTEGER NOT NULL,\n",
    "\t\"InvoiceDate\" DATETIME NOT NULL,\n",
    "\t\"BillingAddress\" NVARCHAR(70),\n",
    "\t\"BillingCity\" NVARCHAR(40),\n",
    "\t\"BillingState\" NVARCHAR(40),\n",
    "\t\"BillingCountry\" NVARCHAR(40),\n",
    "\t\"BillingPostalCode\" NVARCHAR(10),\n",
    "\t\"Total\" NUMERIC(10, 2) NOT NULL,\n",
    "\tPRIMARY KEY (\"InvoiceId\"),\n",
    "\tFOREIGN KEY(\"CustomerId\") REFERENCES \"Customer\" (\"CustomerId\")\n",
    ")\n",
    "\n",
    "/*\n",
    "3 rows from Invoice table:\n",
    "InvoiceId\tCustomerId\tInvoiceDate\tBillingAddress\tBillingCity\tBillingState\tBillingCountry\tBillingPostalCode\tTotal\n",
    "1\t2\t2009-01-01 00:00:00\tTheodor-Heuss-Straße 34\tStuttgart\tNone\tGermany\t70174\t1.98\n",
    "2\t4\t2009-01-02 00:00:00\tUllevålsveien 14\tOslo\tNone\tNorway\t0171\t3.96\n",
    "3\t8\t2009-01-03 00:00:00\tGrétrystraat 63\tBrussels\tNone\tBelgium\t1000\t5.94\n",
    "*/\n",
    "\n",
    "\n",
    "CREATE TABLE \"InvoiceLine\" (\n",
    "\t\"InvoiceLineId\" INTEGER NOT NULL,\n",
    "\t\"InvoiceId\" INTEGER NOT NULL,\n",
    "\t\"TrackId\" INTEGER NOT NULL,\n",
    "\t\"UnitPrice\" NUMERIC(10, 2) NOT NULL,\n",
    "\t\"Quantity\" INTEGER NOT NULL,\n",
    "\tPRIMARY KEY (\"InvoiceLineId\"),\n",
    "\tFOREIGN KEY(\"TrackId\") REFERENCES \"Track\" (\"TrackId\"),\n",
    "\tFOREIGN KEY(\"InvoiceId\") REFERENCES \"Invoice\" (\"InvoiceId\")\n",
    ")\n",
    "\n",
    "/*\n",
    "3 rows from InvoiceLine table:\n",
    "InvoiceLineId\tInvoiceId\tTrackId\tUnitPrice\tQuantity\n",
    "1\t1\t2\t0.99\t1\n",
    "2\t1\t4\t0.99\t1\n",
    "3\t2\t6\t0.99\t1\n",
    "*/\n",
    "\n",
    "\n",
    "CREATE TABLE \"MediaType\" (\n",
    "\t\"MediaTypeId\" INTEGER NOT NULL,\n",
    "\t\"Name\" NVARCHAR(120),\n",
    "\tPRIMARY KEY (\"MediaTypeId\")\n",
    ")\n",
    "\n",
    "/*\n",
    "3 rows from MediaType table:\n",
    "MediaTypeId\tName\n",
    "1\tMPEG audio file\n",
    "2\tProtected AAC audio file\n",
    "3\tProtected MPEG-4 video file\n",
    "*/\n",
    "\n",
    "\n",
    "CREATE TABLE \"Playlist\" (\n",
    "\t\"PlaylistId\" INTEGER NOT NULL,\n",
    "\t\"Name\" NVARCHAR(120),\n",
    "\tPRIMARY KEY (\"PlaylistId\")\n",
    ")\n",
    "\n",
    "/*\n",
    "3 rows from Playlist table:\n",
    "PlaylistId\tName\n",
    "1\tMusic\n",
    "2\tMovies\n",
    "3\tTV Shows\n",
    "*/\n",
    "\n",
    "\n",
    "CREATE TABLE \"PlaylistTrack\" (\n",
    "\t\"PlaylistId\" INTEGER NOT NULL,\n",
    "\t\"TrackId\" INTEGER NOT NULL,\n",
    "\tPRIMARY KEY (\"PlaylistId\", \"TrackId\"),\n",
    "\tFOREIGN KEY(\"TrackId\") REFERENCES \"Track\" (\"TrackId\"),\n",
    "\tFOREIGN KEY(\"PlaylistId\") REFERENCES \"Playlist\" (\"PlaylistId\")\n",
    ")\n",
    "\n",
    "/*\n",
    "3 rows from PlaylistTrack table:\n",
    "PlaylistId\tTrackId\n",
    "1\t3402\n",
    "1\t3389\n",
    "1\t3390\n",
    "*/\n",
    "\n",
    "\n",
    "CREATE TABLE \"Track\" (\n",
    "\t\"TrackId\" INTEGER NOT NULL,\n",
    "\t\"Name\" NVARCHAR(200) NOT NULL,\n",
    "\t\"AlbumId\" INTEGER,\n",
    "\t\"MediaTypeId\" INTEGER NOT NULL,\n",
    "\t\"GenreId\" INTEGER,\n",
    "\t\"Composer\" NVARCHAR(220),\n",
    "\t\"Milliseconds\" INTEGER NOT NULL,\n",
    "\t\"Bytes\" INTEGER,\n",
    "\t\"UnitPrice\" NUMERIC(10, 2) NOT NULL,\n",
    "\tPRIMARY KEY (\"TrackId\"),\n",
    "\tFOREIGN KEY(\"MediaTypeId\") REFERENCES \"MediaType\" (\"MediaTypeId\"),\n",
    "\tFOREIGN KEY(\"GenreId\") REFERENCES \"Genre\" (\"GenreId\"),\n",
    "\tFOREIGN KEY(\"AlbumId\") REFERENCES \"Album\" (\"AlbumId\")\n",
    ")\n",
    "\n",
    "/*\n",
    "3 rows from Track table:\n",
    "TrackId\tName\tAlbumId\tMediaTypeId\tGenreId\tComposer\tMilliseconds\tBytes\tUnitPrice\n",
    "1\tFor Those About To Rock (We Salute You)\t1\t1\t1\tAngus Young, Malcolm Young, Brian Johnson\t343719\t11170334\t0.99\n",
    "2\tBalls to the Wall\t2\t2\t1\tNone\t342562\t5510424\t0.99\n",
    "3\tFast As a Shark\t3\t2\t1\tF. Baltes, S. Kaufman, U. Dirkscneider & W. Hoffman\t230619\t3990994\t0.99\n",
    "*/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13aa3715",
   "metadata": {},
   "source": [
    "<!-- 建表语句-mysql -->\n",
    "-- 艺人表\n",
    "CREATE TABLE `Artist` (\n",
    "  `ArtistId` INT NOT NULL AUTO_INCREMENT,\n",
    "  `Name` VARCHAR(120) DEFAULT NULL,\n",
    "  PRIMARY KEY (`ArtistId`)\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='艺人表';\n",
    "\n",
    "-- 插入艺人表示例数据\n",
    "INSERT INTO `Artist` (`ArtistId`, `Name`) VALUES\n",
    "(1, 'AC/DC'),\n",
    "(2, 'Accept'),\n",
    "(3, 'Aerosmith');\n",
    "\n",
    "-- 专辑表\n",
    "CREATE TABLE `Album` (\n",
    "  `AlbumId` INT NOT NULL AUTO_INCREMENT,\n",
    "  `Title` VARCHAR(160) NOT NULL,\n",
    "  `ArtistId` INT NOT NULL,\n",
    "  PRIMARY KEY (`AlbumId`),\n",
    "  FOREIGN KEY (`ArtistId`) REFERENCES `Artist` (`ArtistId`) ON DELETE RESTRICT ON UPDATE CASCADE\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='专辑表';\n",
    "\n",
    "-- 插入专辑表示例数据\n",
    "INSERT INTO `Album` (`AlbumId`, `Title`, `ArtistId`) VALUES\n",
    "(1, 'For Those About To Rock We Salute You', 1),\n",
    "(2, 'Balls to the Wall', 2),\n",
    "(3, 'Restless and Wild', 2);\n",
    "\n",
    "-- 员工表\n",
    "CREATE TABLE `Employee` (\n",
    "  `EmployeeId` INT NOT NULL AUTO_INCREMENT,\n",
    "  `LastName` VARCHAR(20) NOT NULL,\n",
    "  `FirstName` VARCHAR(20) NOT NULL,\n",
    "  `Title` VARCHAR(30) DEFAULT NULL,\n",
    "  `ReportsTo` INT DEFAULT NULL,\n",
    "  `BirthDate` DATETIME DEFAULT NULL,\n",
    "  `HireDate` DATETIME DEFAULT NULL,\n",
    "  `Address` VARCHAR(70) DEFAULT NULL,\n",
    "  `City` VARCHAR(40) DEFAULT NULL,\n",
    "  `State` VARCHAR(40) DEFAULT NULL,\n",
    "  `Country` VARCHAR(40) DEFAULT NULL,\n",
    "  `PostalCode` VARCHAR(10) DEFAULT NULL,\n",
    "  `Phone` VARCHAR(24) DEFAULT NULL,\n",
    "  `Fax` VARCHAR(24) DEFAULT NULL,\n",
    "  `Email` VARCHAR(60) DEFAULT NULL,\n",
    "  PRIMARY KEY (`EmployeeId`),\n",
    "  FOREIGN KEY (`ReportsTo`) REFERENCES `Employee` (`EmployeeId`) ON DELETE RESTRICT ON UPDATE CASCADE\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='员工表';\n",
    "\n",
    "-- 插入员工表示例数据\n",
    "INSERT INTO `Employee` (`EmployeeId`, `LastName`, `FirstName`, `Title`, `ReportsTo`, `BirthDate`, `HireDate`, `Address`, `City`, `State`, `Country`, `PostalCode`, `Phone`, `Fax`, `Email`) VALUES\n",
    "(1, 'Adams', 'Andrew', 'General Manager', NULL, '1962-02-18 00:00:00', '2002-08-14 00:00:00', '11120 Jasper Ave NW', 'Edmonton', 'AB', 'Canada', 'T5K 2N1', '+1 (780) 428-9482', '+1 (780) 428-3457', 'andrew@chinookcorp.com'),\n",
    "(2, 'Edwards', 'Nancy', 'Sales Manager', 1, '1958-12-08 00:00:00', '2002-05-01 00:00:00', '825 8 Ave SW', 'Calgary', 'AB', 'Canada', 'T2P 2T3', '+1 (403) 262-3443', '+1 (403) 262-3322', 'nancy@chinookcorp.com'),\n",
    "(3, 'Peacock', 'Jane', 'Sales Support Agent', 2, '1973-08-29 00:00:00', '2002-04-01 00:00:00', '1111 6 Ave SW', 'Calgary', 'AB', 'Canada', 'T2P 5M5', '+1 (403) 262-3443', '+1 (403) 262-6712', 'jane@chinookcorp.com');\n",
    "\n",
    "-- 客户表\n",
    "CREATE TABLE `Customer` (\n",
    "  `CustomerId` INT NOT NULL AUTO_INCREMENT,\n",
    "  `FirstName` VARCHAR(40) NOT NULL,\n",
    "  `LastName` VARCHAR(20) NOT NULL,\n",
    "  `Company` VARCHAR(80) DEFAULT NULL,\n",
    "  `Address` VARCHAR(70) DEFAULT NULL,\n",
    "  `City` VARCHAR(40) DEFAULT NULL,\n",
    "  `State` VARCHAR(40) DEFAULT NULL,\n",
    "  `Country` VARCHAR(40) DEFAULT NULL,\n",
    "  `PostalCode` VARCHAR(10) DEFAULT NULL,\n",
    "  `Phone` VARCHAR(24) DEFAULT NULL,\n",
    "  `Fax` VARCHAR(24) DEFAULT NULL,\n",
    "  `Email` VARCHAR(60) NOT NULL,\n",
    "  `SupportRepId` INT DEFAULT NULL,\n",
    "  PRIMARY KEY (`CustomerId`),\n",
    "  FOREIGN KEY (`SupportRepId`) REFERENCES `Employee` (`EmployeeId`) ON DELETE RESTRICT ON UPDATE CASCADE\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='客户表';\n",
    "\n",
    "-- 插入客户表示例数据\n",
    "INSERT INTO `Customer` (`CustomerId`, `FirstName`, `LastName`, `Company`, `Address`, `City`, `State`, `Country`, `PostalCode`, `Phone`, `Fax`, `Email`, `SupportRepId`) VALUES\n",
    "(1, 'Luís', 'Gonçalves', 'Embraer - Empresa Brasileira de Aeronáutica S.A.', 'Av. Brigadeiro Faria Lima, 2170', 'São José dos Campos', 'SP', 'Brazil', '12227-000', '+55 (12) 3923-5555', '+55 (12) 3923-5566', 'luisg@embraer.com.br', 3),\n",
    "(2, 'Leonie', 'Köhler', NULL, 'Theodor-Heuss-Straße 34', 'Stuttgart', NULL, 'Germany', '70174', '+49 0711 2842222', NULL, 'leonekohler@surfeu.de', 5),\n",
    "(3, 'François', 'Tremblay', NULL, '1498 rue Bélanger', 'Montréal', 'QC', 'Canada', 'H2G 1A7', '+1 (514) 721-4711', NULL, 'ftremblay@gmail.com', 3);\n",
    "\n",
    "-- 流派表\n",
    "CREATE TABLE `Genre` (\n",
    "  `GenreId` INT NOT NULL AUTO_INCREMENT,\n",
    "  `Name` VARCHAR(120) DEFAULT NULL,\n",
    "  PRIMARY KEY (`GenreId`)\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='流派表';\n",
    "\n",
    "-- 插入流派表示例数据\n",
    "INSERT INTO `Genre` (`GenreId`, `Name`) VALUES\n",
    "(1, 'Rock'),\n",
    "(2, 'Jazz'),\n",
    "(3, 'Metal');\n",
    "\n",
    "-- 媒体类型表\n",
    "CREATE TABLE `MediaType` (\n",
    "  `MediaTypeId` INT NOT NULL AUTO_INCREMENT,\n",
    "  `Name` VARCHAR(120) DEFAULT NULL,\n",
    "  PRIMARY KEY (`MediaTypeId`)\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='媒体类型表';\n",
    "\n",
    "-- 插入媒体类型表示例数据\n",
    "INSERT INTO `MediaType` (`MediaTypeId`, `Name`) VALUES\n",
    "(1, 'MPEG audio file'),\n",
    "(2, 'Protected AAC audio file'),\n",
    "(3, 'Protected MPEG-4 video file');\n",
    "\n",
    "-- 曲目表\n",
    "CREATE TABLE `Track` (\n",
    "  `TrackId` INT NOT NULL AUTO_INCREMENT,\n",
    "  `Name` VARCHAR(200) NOT NULL,\n",
    "  `AlbumId` INT DEFAULT NULL,\n",
    "  `MediaTypeId` INT NOT NULL,\n",
    "  `GenreId` INT DEFAULT NULL,\n",
    "  `Composer` VARCHAR(220) DEFAULT NULL,\n",
    "  `Milliseconds` INT NOT NULL,\n",
    "  `Bytes` INT DEFAULT NULL,\n",
    "  `UnitPrice` DECIMAL(10, 2) NOT NULL,\n",
    "  PRIMARY KEY (`TrackId`),\n",
    "  FOREIGN KEY (`MediaTypeId`) REFERENCES `MediaType` (`MediaTypeId`) ON DELETE RESTRICT ON UPDATE CASCADE,\n",
    "  FOREIGN KEY (`GenreId`) REFERENCES `Genre` (`GenreId`) ON DELETE RESTRICT ON UPDATE CASCADE,\n",
    "  FOREIGN KEY (`AlbumId`) REFERENCES `Album` (`AlbumId`) ON DELETE RESTRICT ON UPDATE CASCADE\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='曲目表';\n",
    "\n",
    "-- 插入曲目表示例数据\n",
    "INSERT INTO `Track` (`TrackId`, `Name`, `AlbumId`, `MediaTypeId`, `GenreId`, `Composer`, `Milliseconds`, `Bytes`, `UnitPrice`) VALUES\n",
    "(1, 'For Those About To Rock (We Salute You)', 1, 1, 1, 'Angus Young, Malcolm Young, Brian Johnson', 343719, 11170334, 0.99),\n",
    "(2, 'Balls to the Wall', 2, 2, 1, NULL, 342562, 5510424, 0.99),\n",
    "(3, 'Fast As a Shark', 3, 2, 1, 'F. Baltes, S. Kaufman, U. Dirkscneider & W. Hoffman', 230619, 3990994, 0.99);\n",
    "\n",
    "-- 播放列表表\n",
    "CREATE TABLE `Playlist` (\n",
    "  `PlaylistId` INT NOT NULL AUTO_INCREMENT,\n",
    "  `Name` VARCHAR(120) DEFAULT NULL,\n",
    "  PRIMARY KEY (`PlaylistId`)\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='播放列表表';\n",
    "\n",
    "-- 插入播放列表表示例数据\n",
    "INSERT INTO `Playlist` (`PlaylistId`, `Name`) VALUES\n",
    "(1, 'Music'),\n",
    "(2, 'Movies'),\n",
    "(3, 'TV Shows');\n",
    "\n",
    "-- 播放列表曲目关联表\n",
    "CREATE TABLE `PlaylistTrack` (\n",
    "  `PlaylistId` INT NOT NULL,\n",
    "  `TrackId` INT NOT NULL,\n",
    "  PRIMARY KEY (`PlaylistId`, `TrackId`),\n",
    "  FOREIGN KEY (`TrackId`) REFERENCES `Track` (`TrackId`) ON DELETE RESTRICT ON UPDATE CASCADE,\n",
    "  FOREIGN KEY (`PlaylistId`) REFERENCES `Playlist` (`PlaylistId`) ON DELETE RESTRICT ON UPDATE CASCADE\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='播放列表曲目关联表';\n",
    "\n",
    "-- 插入播放列表曲目关联表示例数据\n",
    "INSERT INTO `PlaylistTrack` (`PlaylistId`, `TrackId`) VALUES\n",
    "(1, 3402),\n",
    "(1, 3389),\n",
    "(1, 3390);\n",
    "\n",
    "-- 发票表\n",
    "CREATE TABLE `Invoice` (\n",
    "  `InvoiceId` INT NOT NULL AUTO_INCREMENT,\n",
    "  `CustomerId` INT NOT NULL,\n",
    "  `InvoiceDate` DATETIME NOT NULL,\n",
    "  `BillingAddress` VARCHAR(70) DEFAULT NULL,\n",
    "  `BillingCity` VARCHAR(40) DEFAULT NULL,\n",
    "  `BillingState` VARCHAR(40) DEFAULT NULL,\n",
    "  `BillingCountry` VARCHAR(40) DEFAULT NULL,\n",
    "  `BillingPostalCode` VARCHAR(10) DEFAULT NULL,\n",
    "  `Total` DECIMAL(10, 2) NOT NULL,\n",
    "  PRIMARY KEY (`InvoiceId`),\n",
    "  FOREIGN KEY (`CustomerId`) REFERENCES `Customer` (`CustomerId`) ON DELETE RESTRICT ON UPDATE CASCADE\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='发票表';\n",
    "\n",
    "-- 插入发票表示例数据\n",
    "INSERT INTO `Invoice` (`InvoiceId`, `CustomerId`, `InvoiceDate`, `BillingAddress`, `BillingCity`, `BillingState`, `BillingCountry`, `BillingPostalCode`, `Total`) VALUES\n",
    "(1, 2, '2009-01-01 00:00:00', 'Theodor-Heuss-Straße 34', 'Stuttgart', NULL, 'Germany', '70174', 1.98),\n",
    "(2, 4, '2009-01-02 00:00:00', 'Ullevålsveien 14', 'Oslo', NULL, 'Norway', '0171', 3.96),\n",
    "(3, 8, '2009-01-03 00:00:00', 'Grétrystraat 63', 'Brussels', NULL, 'Belgium', '1000', 5.94);\n",
    "\n",
    "-- 发票明细表\n",
    "CREATE TABLE `InvoiceLine` (\n",
    "  `InvoiceLineId` INT NOT NULL AUTO_INCREMENT,\n",
    "  `InvoiceId` INT NOT NULL,\n",
    "  `TrackId` INT NOT NULL,\n",
    "  `UnitPrice` DECIMAL(10, 2) NOT NULL,\n",
    "  `Quantity` INT NOT NULL,\n",
    "  PRIMARY KEY (`InvoiceLineId`),\n",
    "  FOREIGN KEY (`TrackId`) REFERENCES `Track` (`TrackId`) ON DELETE RESTRICT ON UPDATE CASCADE,\n",
    "  FOREIGN KEY (`InvoiceId`) REFERENCES `Invoice` (`InvoiceId`) ON DELETE RESTRICT ON UPDATE CASCADE\n",
    ") ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COMMENT='发票明细表';\n",
    "\n",
    "-- 插入发票明细表示例数据\n",
    "INSERT INTO `InvoiceLine` (`InvoiceLineId`, `InvoiceId`, `TrackId`, `UnitPrice`, `Quantity`) VALUES\n",
    "(1, 1, 2, 0.99, 1),\n",
    "(2, 1, 4, 0.99, 1),\n",
    "(3, 2, 6, 0.99, 1);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ac9cecf",
   "metadata": {},
   "outputs": [],
   "source": [
    "SCHEMA = db.get_table_info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bdf9ca5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "from langchain_core.tools import tool\n",
    "DENY_RE = re.compile(r\"\\b(INSERT|UPDATE|DELETE|ALTER|DROP|CREATE|REPLACE|TRUNCATE)\\b\", re.I)\n",
    "HAS_LIMIT_TAIL_RE = re.compile(r\"(?is)\\blimit\\b\\s+\\d+(\\s*,\\s*\\d+)?\\s*;?\\s*$\")\n",
    "\n",
    "def _safe_sql(q: str) -> str:\n",
    "    # normalize\n",
    "    q = q.strip()\n",
    "    # block multiple statements (allow one optional trailing ;)\n",
    "    if q.count(\";\") > 1 or (q.endswith(\";\") and \";\" in q[:-1]):\n",
    "        return \"Error: multiple statements are not allowed.\"\n",
    "    q = q.rstrip(\";\").strip()\n",
    "\n",
    "    # read-only gate\n",
    "    if not q.lower().startswith(\"select\"):\n",
    "        return \"Error: only SELECT statements are allowed.\"\n",
    "    if DENY_RE.search(q):\n",
    "        return \"Error: DML/DDL detected. Only read-only queries are permitted.\"\n",
    "\n",
    "    # append LIMIT only if not already present at the end (robust to whitespace/newlines)\n",
    "    if not HAS_LIMIT_TAIL_RE.search(q):\n",
    "        q += \" LIMIT 5\"\n",
    "    return q\n",
    "\n",
    "@tool\n",
    "def execute_sql(query: str) -> str:\n",
    "    \"\"\"Execute a READ-ONLY SQLite SELECT query and return results.\"\"\"\n",
    "    query = _safe_sql(query)\n",
    "    q = query\n",
    "    if q.startswith(\"Error:\"):\n",
    "        return q\n",
    "    try:\n",
    "        return db.run(q)\n",
    "    except Exception as e:\n",
    "        return f\"Error: {e}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da3377b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM = f\"\"\"You are a careful SQLite analyst.\n",
    "\n",
    "Authoritative schema (do not invent columns/tables):\n",
    "{SCHEMA}\n",
    "\n",
    "Rules:\n",
    "- Think step-by-step.\n",
    "- When you need data, call the tool `execute_sql` with ONE SELECT query.\n",
    "- Read-only only; no INSERT/UPDATE/DELETE/ALTER/DROP/CREATE/REPLACE/TRUNCATE.\n",
    "- Limit to 5 rows unless user explicitly asks otherwise.\n",
    "- If the tool returns 'Error:', revise the SQL and try again.\n",
    "- Limit the number of attempts to 5.\n",
    "- If you are not successful after 5 attempts, return a note to the user.\n",
    "- Prefer explicit column lists; avoid SELECT *.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6d3afb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import create_agent\n",
    "from langchain_core.messages import SystemMessage\n",
    "agent = create_agent(\n",
    "    model=llm,\n",
    "    tools=[execute_sql],\n",
    "    prompt=SystemMessage(content=SYSTEM),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02e35612",
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Which genre on average has the longest tracks?\"\n",
    "\n",
    "for step in agent.stream(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": question}]},\n",
    "    stream_mode=\"values\",\n",
    "):\n",
    "    step[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c6481cb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pymysql\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "\n",
    "# 连接到MySQL数据库（请替换为您的实际连接信息）\n",
    "db = SQLDatabase.from_uri(\"mysql+pymysql://root:root@127.0.0.1:3306/langchain_test\")\n",
    "\n",
    "SCHEMA = db.get_table_info()\n",
    "\n",
    "import re\n",
    "from langchain_core.tools import tool\n",
    "DENY_RE = re.compile(r\"\\b(INSERT|UPDATE|DELETE|ALTER|DROP|CREATE|REPLACE|TRUNCATE)\\b\", re.I)\n",
    "HAS_LIMIT_TAIL_RE = re.compile(r\"(?is)\\blimit\\b\\s+\\d+(\\s*,\\s*\\d+)?\\s*;?\\s*$\")\n",
    "\n",
    "def _safe_sql(q: str) -> str:\n",
    "    # 规范化查询\n",
    "    q = q.strip()\n",
    "    # 阻止多个语句（允许一个可选的尾随;）\n",
    "    if q.count(\";\") > 1 or (q.endswith(\";\") and \";\" in q[:-1]):\n",
    "        return \"错误：不允许多个语句。\"\n",
    "    q = q.rstrip(\";\").strip()\n",
    "\n",
    "    # 只读检查\n",
    "    if not q.lower().startswith(\"select\"):\n",
    "        return \"错误：只允许SELECT语句。\"\n",
    "    if DENY_RE.search(q):\n",
    "        return \"错误：检测到DML/DDL操作。只允许只读查询。\"\n",
    "\n",
    "    # 如果末尾没有LIMIT则添加（对空格/换行符鲁棒）\n",
    "    if not HAS_LIMIT_TAIL_RE.search(q):\n",
    "        q += \" LIMIT 5\"\n",
    "    return q\n",
    "\n",
    "@tool\n",
    "def execute_sql(query: str) -> str:\n",
    "    \"\"\"执行只读的MySQL SELECT查询并返回结果\"\"\"\n",
    "    query = _safe_sql(query)\n",
    "    q = query\n",
    "    if q.startswith(\"错误：\"):\n",
    "        return q\n",
    "    try:\n",
    "        return db.run(q)\n",
    "    except Exception as e:\n",
    "        return f\"错误：{e}\"\n",
    "\n",
    "SYSTEM = f\"\"\"您是一位谨慎的MySQL分析师。\n",
    "\n",
    "数据库架构（请勿虚构列/表）：\n",
    "{SCHEMA}\n",
    "\n",
    "规则：\n",
    "- 逐步思考。\n",
    "- 当需要数据时，使用一个SELECT查询调用工具`execute_sql`。\n",
    "- 仅限只读操作；不允许INSERT/UPDATE/DELETE/ALTER/DROP/CREATE/REPLACE/TRUNCATE。\n",
    "- 除非用户明确要求，否则限制为5行。\n",
    "- 如果工具返回'错误：'，请修改SQL并重试。\n",
    "- 尝试次数限制为5次。\n",
    "- 如果5次尝试后仍未成功，请向用户返回说明。\n",
    "- 优先使用明确的列列表；避免使用SELECT *。\n",
    "\"\"\"\n",
    "\n",
    "from langchain.agents import create_agent\n",
    "from langchain_core.messages import SystemMessage\n",
    "agent = create_agent(\n",
    "    model=llm,\n",
    "    tools=[execute_sql],\n",
    "    prompt=SystemMessage(content=SYSTEM),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b06fb60b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "员工表有多少人？\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "我来查询员工表的总人数。\n",
      "Tool Calls:\n",
      "  execute_sql (call_00_5LyjNGnrLg3peHofHjZqrMAu)\n",
      " Call ID: call_00_5LyjNGnrLg3peHofHjZqrMAu\n",
      "  Args:\n",
      "    query: SELECT COUNT(*) as total_employees FROM employees\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: execute_sql\n",
      "\n",
      "[(7,)]\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "员工表中共有7名员工。\n"
     ]
    }
   ],
   "source": [
    "question = \"员工表有多少人？\"\n",
    "\n",
    "for step in agent.stream(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": question}]},\n",
    "    stream_mode=\"values\",\n",
    "):\n",
    "    step[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d16a2591",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "各个部门中员工工资最高的分别是谁？\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "我需要查询各个部门中工资最高的员工信息。让我逐步分析这个问题。\n",
      "\n",
      "首先，我需要了解部门信息和员工工资数据，然后找到每个部门的最高工资员工。\n",
      "Tool Calls:\n",
      "  execute_sql (call_00_CA9XxVmmtW63lyOCHaYFvfB6)\n",
      " Call ID: call_00_CA9XxVmmtW63lyOCHaYFvfB6\n",
      "  Args:\n",
      "    query: SELECT d.dept_id, d.dept_name, MAX(e.salary) as max_salary\n",
      "FROM departments d\n",
      "JOIN emp_dept_relation edr ON d.dept_id = edr.dept_id\n",
      "JOIN employees e ON edr.emp_id = e.emp_id\n",
      "GROUP BY d.dept_id, d.dept_name\n",
      "ORDER BY d.dept_id\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: execute_sql\n",
      "\n",
      "[(1, '技术部', Decimal('25000.00')), (2, '市场部', Decimal('22000.00')), (3, '人力资源部', Decimal('12000.00')), (4, '财务部', Decimal('18000.00'))]\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "现在我知道了每个部门的最高工资。接下来，我需要找到对应这些最高工资的具体员工信息：\n",
      "Tool Calls:\n",
      "  execute_sql (call_00_u8rGTQ9dNed8nx3GxRbOg65H)\n",
      " Call ID: call_00_u8rGTQ9dNed8nx3GxRbOg65H\n",
      "  Args:\n",
      "    query: SELECT d.dept_id, d.dept_name, e.emp_name, e.position, e.salary\n",
      "FROM departments d\n",
      "JOIN emp_dept_relation edr ON d.dept_id = edr.dept_id\n",
      "JOIN employees e ON edr.emp_id = e.emp_id\n",
      "WHERE (d.dept_id, e.salary) IN (\n",
      "    SELECT d2.dept_id, MAX(e2.salary)\n",
      "    FROM departments d2\n",
      "    JOIN emp_dept_relation edr2 ON d2.dept_id = edr2.dept_id\n",
      "    JOIN employees e2 ON edr2.emp_id = e2.emp_id\n",
      "    GROUP BY d2.dept_id\n",
      ")\n",
      "ORDER BY d.dept_id\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: execute_sql\n",
      "\n",
      "[(1, '技术部', '张三', '高级工程师', Decimal('25000.00')), (2, '市场部', '李四', '产品经理', Decimal('22000.00')), (3, '人力资源部', '孙七', '人事专员', Decimal('12000.00')), (4, '财务部', '赵六', '财务主管', Decimal('18000.00'))]\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "根据查询结果，各个部门中工资最高的员工分别是：\n",
      "\n",
      "1. **技术部**：张三（高级工程师），工资：25,000.00元\n",
      "2. **市场部**：李四（产品经理），工资：22,000.00元  \n",
      "3. **人力资源部**：孙七（人事专员），工资：12,000.00元\n",
      "4. **财务部**：赵六（财务主管），工资：18,000.00元\n",
      "\n",
      "这些是每个部门中薪资最高的员工信息。\n"
     ]
    }
   ],
   "source": [
    "question = \"各个部门中员工工资最高的分别是谁？\"\n",
    "\n",
    "for step in agent.stream(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": question}]},\n",
    "    stream_mode=\"values\",\n",
    "):\n",
    "    step[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e271592",
   "metadata": {},
   "source": [
    "二. Build a customized workflow using LangGraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a5e6c4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize an LLM\n",
    "from langchain.chat_models import init_chat_model\n",
    "\n",
    "llm = init_chat_model(\"deepseek:deepseek-chat\")\n",
    "\n",
    "import pathlib\n",
    "import requests\n",
    "\n",
    "# Initialize the database\n",
    "\n",
    "url = \"https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db\"\n",
    "local_path = pathlib.Path(\"Chinook.db\")\n",
    "\n",
    "if local_path.exists():\n",
    "    print(f\"{local_path} already exists, skipping download.\")\n",
    "else:\n",
    "    response = requests.get(url)\n",
    "    if response.status_code == 200:\n",
    "        local_path.write_bytes(response.content)\n",
    "        print(f\"File downloaded and saved as {local_path}\")\n",
    "    else:\n",
    "        print(f\"Failed to download the file. Status code: {response.status_code}\")\n",
    "\n",
    "db = SQLDatabase.from_uri(\"sqlite:///Chinook.db\")\n",
    "SCHEMA = db.get_table_info()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cba4772d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Graph State\n",
    "class GraphState(MessagesState):\n",
    "    first_name: Optional[str]\n",
    "    last_name: Optional[str]\n",
    "    customer: bool\n",
    "    customer_id: Optional[int]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6cadbe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Policy configuration ------------------------------------------------------\n",
    "\n",
    "# Tables a customer is allowed to read\n",
    "CUSTOMER_ALLOWLIST = {\n",
    "    \"invoice\",\n",
    "    \"invoiceline\",\n",
    "    \"track\",\n",
    "    \"album\",\n",
    "    \"artist\",\n",
    "    \"genre\",\n",
    "    \"mediatype\",\n",
    "    \"playlist\",\n",
    "    \"playlisttrack\",\n",
    "}\n",
    "\n",
    "# Tables that are customer-scoped (must include CustomerId = :customer_id)\n",
    "CUSTOMER_SCOPED = {\"invoice\", \"invoiceline\"}\n",
    "\n",
    "# --- Safety regexes ------------------------------------------------------------\n",
    "DENY_RE = re.compile(r\"\\b(INSERT|UPDATE|DELETE|ALTER|DROP|CREATE|REPLACE|TRUNCATE)\\b\", re.I)\n",
    "HAS_LIMIT_TAIL_RE = re.compile(r\"(?is)\\blimit\\b\\s+\\d+(\\s*,\\s*\\d+)?\\s*;?\\s*$\")\n",
    "\n",
    "# Disallow non-plain-select constructs to keep verification simple\n",
    "NON_PLAIN_SQL_RE = re.compile(r\"\\b(with|union|intersect|except)\\b|\\(\\s*select\\b\", re.I)\n",
    "\n",
    "# Extract FROM/JOIN tables & aliases (very lightweight parsing)\n",
    "FROM_RE = re.compile(r\"\\bfrom\\s+([\\\"`\\[]?\\w+[\\\"`\\]]?)(?:\\s+as\\s+(\\w+)|\\s+(\\w+))?\", re.I)\n",
    "JOIN_RE = re.compile(r\"\\bjoin\\s+([\\\"`\\[]?\\w+[\\\"`\\]]?)(?:\\s+as\\s+(\\w+)|\\s+(\\w+))?\", re.I)\n",
    "\n",
    "# Simple checks around CustomerId usage\n",
    "CUSTID_PLACEHOLDER_EQ_RE = re.compile(r\"\\b(?:\\w+\\.)?customerid\\s*=\\s*:customer_id\\b\", re.I)\n",
    "CUSTID_NUMERIC_EQ_RE     = re.compile(r\"\\b(?:\\w+\\.)?customerid\\s*=\\s*\\d+\\b\", re.I)\n",
    "\n",
    "\n",
    "def _normalize_ident(name: str) -> str:\n",
    "    # strip quotes/backticks/brackets and lower-case\n",
    "    return re.sub(r'^[\\\"`\\[]|[\\\"`\\]]$', '', name).lower()\n",
    "\n",
    "\n",
    "def _extract_tables_and_aliases(q: str):\n",
    "    tables = set()\n",
    "    alias_map = {}  # alias -> base table (lower-cased)\n",
    "    for m in FROM_RE.finditer(q):\n",
    "        base = _normalize_ident(m.group(1))\n",
    "        alias = (m.group(2) or m.group(3) or \"\").lower()\n",
    "        tables.add(base)\n",
    "        if alias:\n",
    "            alias_map[alias] = base\n",
    "    for m in JOIN_RE.finditer(q):\n",
    "        base = _normalize_ident(m.group(1))\n",
    "        alias = (m.group(2) or m.group(3) or \"\").lower()\n",
    "        tables.add(base)\n",
    "        if alias:\n",
    "            alias_map[alias] = base\n",
    "    return tables, alias_map\n",
    "\n",
    "\n",
    "def _safe_sql(q: str, customer_id: int) -> str:\n",
    "    # normalize\n",
    "    q = q.strip()\n",
    "    # block multiple statements (allow one optional trailing ;)\n",
    "    if q.count(\";\") > 1 or (q.endswith(\";\") and \";\" in q[:-1]):\n",
    "        return \"Error: multiple statements are not allowed.\"\n",
    "    q = q.rstrip(\";\").strip()\n",
    "\n",
    "    # read-only gate\n",
    "    if not q.lower().startswith(\"select\"):\n",
    "        return \"Error: only SELECT statements are allowed.\"\n",
    "    if DENY_RE.search(q):\n",
    "        return \"Error: DML/DDL detected. Only read-only queries are permitted.\"\n",
    "\n",
    "    # plain-select only (no CTEs, subqueries, UNION/INTERSECT/EXCEPT)\n",
    "    if NON_PLAIN_SQL_RE.search(q):\n",
    "        return \"Error: only plain SELECTs (no CTEs/subqueries/UNION/INTERSECT/EXCEPT) are allowed.\"\n",
    "\n",
    "    # gather referenced tables & aliases\n",
    "    tables, alias_map = _extract_tables_and_aliases(q)\n",
    "    if not tables:\n",
    "        return \"Error: could not determine referenced tables.\"\n",
    "\n",
    "    # allowlist enforcement\n",
    "    disallowed = {t for t in tables if t not in CUSTOMER_ALLOWLIST}\n",
    "    if disallowed:\n",
    "        bad = \", \".join(sorted(disallowed))\n",
    "        return f\"Error: access to tables [{bad}] is not permitted.\"\n",
    "\n",
    "    # customer-scoped enforcement\n",
    "    needs_customer_filter = bool(CUSTOMER_SCOPED & tables)\n",
    "    if needs_customer_filter:\n",
    "        # forbid numeric literals for CustomerId\n",
    "        if CUSTID_NUMERIC_EQ_RE.search(q):\n",
    "            return \"Error: use the :customer_id placeholder (no numeric literals) for CustomerId.\"\n",
    "\n",
    "        # require a CustomerId = :customer_id predicate in the query text\n",
    "        if not CUSTID_PLACEHOLDER_EQ_RE.search(q):\n",
    "            return \"Error: queries touching Invoice/InvoiceLine must include CustomerId = :customer_id.\"\n",
    "\n",
    "        # Special rule for InvoiceLine: must also reference Invoice (joined)\n",
    "        if \"invoiceline\" in tables and \"invoice\" not in tables:\n",
    "            return \"Error: queries referencing InvoiceLine must also join Invoice and filter by CustomerId = :customer_id.\"\n",
    "\n",
    "    # append LIMIT if not present at the end (robust to whitespace/newlines)\n",
    "    if not HAS_LIMIT_TAIL_RE.search(q):\n",
    "        q += \" LIMIT 5\"\n",
    "    return q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8566a0b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool(parse_docstring=True)\n",
    "def execute_sql(\n",
    "    query: str,\n",
    "    state: Annotated[GraphState, InjectedState],  # provides access to customer_id\n",
    ") -> str:\n",
    "    \"\"\"Execute a READ-ONLY SQLite SELECT query (customer-scoped) and return results.\n",
    "\n",
    "    Args:\n",
    "        query: a string containing a valid SQL query\n",
    "\n",
    "    Returns:\n",
    "        A string with the response to the query or an error\n",
    "    \"\"\"\n",
    "    customer_id = int(state[\"customer_id\"])\n",
    "    safe_q = _safe_sql(query, customer_id)\n",
    "    if safe_q.startswith(\"Error:\"):\n",
    "        return safe_q\n",
    "    try:\n",
    "        # Bind the named parameter expected by the query (:customer_id)\n",
    "        return db.run(safe_q, parameters={\"customer_id\": customer_id})\n",
    "    except Exception as e:\n",
    "        return f\"Error: {e}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c465282",
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM = \"\"\"You are a careful SQLite analyst.\n",
    "\n",
    "Authoritative schema (do not invent columns/tables):\n",
    "{SCHEMA}\n",
    "\n",
    "Always use the `:customer_id` placeholder; never hardcode IDs or use names.\n",
    "The system binds the actual value at execution.\n",
    "\n",
    "Rules:\n",
    "- Think step-by-step.\n",
    "- When you need data, call the tool `execute_sql` with ONE SELECT query.\n",
    "- Read-only only; no INSERT/UPDATE/DELETE/ALTER/DROP/CREATE/REPLACE/TRUNCATE.\n",
    "- Limit to 5 rows unless the user explicitly asks otherwise.\n",
    "- If the tool returns 'Error:', revise the SQL and try again.\n",
    "- Limit the number of attempts to 5.\n",
    "- If you are not successful after 5 attempts, return a note to the user.\n",
    "- Prefer explicit column lists; avoid SELECT *.\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30c7ccbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "\n",
    "_ID_RE = re.compile(r\"\\b\\d+\\b\")  # first integer in the run() string\n",
    "\n",
    "def identify_node(state: GraphState) -> GraphState:\n",
    "    first = (state.get(\"first_name\") or \"\").strip()\n",
    "    last  = (state.get(\"last_name\") or \"\").strip()\n",
    "\n",
    "    if not (first and last):\n",
    "        return {}  # nothing to change\n",
    "\n",
    "    # simple quote escaping for SQL string literal\n",
    "    sf = first.replace(\"'\", \"''\")\n",
    "    sl = last.replace(\"'\", \"''\")\n",
    "\n",
    "    try:\n",
    "        cust_raw = db.run(\n",
    "            \"SELECT CustomerId FROM Customer \"\n",
    "            f\"WHERE FirstName = '{sf}' AND LastName = '{sl}' \"\n",
    "            \"LIMIT 1\"\n",
    "        )\n",
    "        if not cust_raw:\n",
    "            return {}  # no change\n",
    "\n",
    "        m = _ID_RE.search(cust_raw)\n",
    "        if not m:\n",
    "            # couldn't parse an ID; don't crash—just no update\n",
    "            return {}\n",
    "\n",
    "        customer_id = int(m.group(0))\n",
    "        return {\n",
    "            \"customer\": True,\n",
    "            \"customer_id\": customer_id,\n",
    "        }\n",
    "\n",
    "    except Exception as e:\n",
    "        print(f\"Customer lookup failed: {e}\")\n",
    "        return {}\n",
    "\n",
    "# conditional edge\n",
    "def route_from_identify(state: GraphState):\n",
    "    # Continue only if an ID is present; otherwise END\n",
    "    if state.get(\"employee_id\") or state.get(\"customer_id\"):\n",
    "        return \"llm\"\n",
    "    return \"unknown_user\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c140ab1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Node Return Unknown User Message\n",
    "def unknown_user_node(state: GraphState):\n",
    "    return {\n",
    "        \"messages\": AIMessage(\n",
    "            f\"The user, first_name:{state.get('first_name','missing')}, \"\n",
    "            f\"last_name:{state.get('last_name','missing')} is not in the database\"\n",
    "        )\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be1999d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Node LLM ReAct step\n",
    "model_with_tools = llm.bind_tools([execute_sql])\n",
    "\n",
    "def llm_node(state: GraphState) -> GraphState:\n",
    "    msgs = [SystemMessage(content=SYSTEM.format(SCHEMA=SCHEMA))] + state[\"messages\"]\n",
    "    ai: AIMessage = model_with_tools.invoke(msgs)\n",
    "    return { \"messages\": [ai]}\n",
    "\n",
    "def route_from_llm(state: GraphState):\n",
    "    last = state[\"messages\"][-1]\n",
    "    if isinstance(last, AIMessage) and getattr(last, \"tool_calls\", None):\n",
    "        return \"tools\"\n",
    "    return END\n",
    "\n",
    "\n",
    "# Node : Tool execution\n",
    "tool_node = ToolNode([execute_sql])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdc8aaf1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build Graph\n",
    "builder = StateGraph(GraphState)\n",
    "\n",
    "builder.add_node(\"identify\", identify_node)\n",
    "builder.add_node(\"unknown_user\", unknown_user_node)\n",
    "builder.add_node(\"llm\", llm_node)\n",
    "builder.add_node(\"tools\", tool_node)\n",
    "\n",
    "builder.set_entry_point(\"identify\")\n",
    "builder.add_conditional_edges(\"identify\", route_from_identify, {\"llm\": \"llm\", \"unknown_user\": \"unknown_user\"})\n",
    "builder.add_conditional_edges(\"llm\", route_from_llm, {\"tools\": \"tools\", END: END})\n",
    "builder.add_edge(\"tools\", \"llm\")\n",
    "\n",
    "graph = builder.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f2ec3fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image, display\n",
    "from langchain_core.runnables.graph import CurveStyle, MermaidDrawMethod, NodeStyles\n",
    "\n",
    "display(Image(graph.get_graph().draw_mermaid_png()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "876da47f",
   "metadata": {},
   "outputs": [],
   "source": [
    "question = \"Show me my last 3 invoices.\"\n",
    "for step in graph.stream(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": question}],\n",
    "      \"first_name\": \"Frank\",\n",
    "      \"last_name\": \"Harris\",\n",
    "    },\n",
    "    stream_mode=\"values\",\n",
    "):\n",
    "    step[\"messages\"][-1].pretty_print()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain_learn",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
